# 通过一个最简单的神经网线，实现二分类任务
import torch
import torch.nn as nn
import matplotlib.pyplot as plt

# 1、定义输入、输出、隐藏层大小、批大小
n_in, n_h, n_out, batch_size = 10, 5, 1, 10

# 2、定义输入数据和标签
x = torch.randn(batch_size, n_in)  # 随机生成输入数据
y = torch.tensor([[1.0], [0.0], [0.0], [1.0], [1.0], [1.0], [0.0], [0.0], [1.0], [1.0]])

# 3、创建模型顺序，包含线性层、ReLU激活函数和Sigmoid激活函数
model = nn.Sequential(
    nn.Linear(n_in, n_h),
    nn.ReLU(),
    nn.Linear(n_h, n_out),
    nn.Sigmoid()
)

# 4、定义均方误差损失函数和随机梯度下降优化器
criterion = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

# 存储每轮损失
losses = []

# 5、训练模型
for epoch in range(50):
    # 前向传播
    y_pred = model(x)
    # 计算损失
    loss = criterion(y_pred, y)
    # 存储损失
    losses.append(loss.item())
    # 打印损失
    print(f'Epoch: {epoch+1}/50, Loss: {loss.item():.4f}')
    # 清空梯度
    optimizer.zero_grad()
    # 反向传播，计算梯度
    loss.backward()
    # 更新参数
    optimizer.step()

# 6、可视化损失
plt.figure(figsize=(8, 5))
plt.plot(range(1, 51), losses, label='Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.title('Training Loss Over Epochs')
plt.legend()
plt.grid()
plt.show()

# 可视化预测结果与实际目标值对比
y_pred_final = model(x).detach().numpy()  # 最终预测值
y_actual = y.numpy()  # 实际值

plt.figure(figsize=(8, 5))
plt.plot(range(1, batch_size + 1), y_actual, 'o-', label='Actual', color='blue')
plt.plot(range(1, batch_size + 1), y_pred_final, 'x--', label='Predicted', color='red')
plt.xlabel('Sample Index')
plt.ylabel('Value')
plt.title('Actual vs Predicted Values')
plt.legend()
plt.grid()
plt.show()